# encoding=utf-8

"""py-file-send.py

    封装usb相关的函数，发送文件到AIPro上

    低速数据格式如下:
        头信息: 空格, int32_1, int32_2, int32_3, int32_4, b_len, finish
            - 空格 uint8: 通信线位丢失修正
            - int32_1_4: 包编号
            - b_len uint8: 后续有效byte个数
            - finish uint8: 是否为最后一个包
        数据体:
            - uint8的数据体: 最大长度(48 - 7)
"""
import time
import traceback

import cv2 as cv
import numpy as np
import usb.core
import usb.util

VID = 0xF055
PID = 0x9999


def init_usb():
    # find our device
    dev = usb.core.find(idVendor=VID, idProduct=PID)

    # was it found?
    if dev is None:
        raise ValueError('Device not found')

    # get an endpoint instance
    cfg = dev.get_active_configuration()
    intf = cfg[(0, 0)]

    outep = usb.util.find_descriptor(
        intf,
        # match the first OUT endpoint
        custom_match= \
            lambda e: \
                usb.util.endpoint_direction(e.bEndpointAddress) == \
                usb.util.ENDPOINT_OUT)

    inep = usb.util.find_descriptor(
        intf,
        # match the first IN endpoint
        custom_match= \
            lambda e: \
                usb.util.endpoint_direction(e.bEndpointAddress) == \
                usb.util.ENDPOINT_IN)

    assert inep is not None
    assert outep is not None

    return outep, inep


def send_raw(outep, data):
    try:
        outep.write(data)
        return True
    except Exception as e:
        print(traceback.format_exc())
        return False


def read_raw(inep, data_len):
    try:
        return inep.read(data_len)
    except Exception as e:
        return None


def send_and_read(outep, inep, data):
    print("next will be wirte: {}, len: {}".format(data, len(data)))
    # clean cache
    read_raw(inep, len(data))
    # send
    if not send_raw(outep, data):
        return False
    # clean cache
    read_raw(inep, len(data))
    # status
    return True


def make_one_frame(buffer: np.ndarray, package_idx: int, start_idx: int, end_idx: int, is_finish: int):
    data_frame = bytearray()
    # 第1个字节为空
    # data_frame = [0x00]
    data_frame.append(0x00)
    # 2-5: 包编号
    for i in range(4):
        uint32_package = (package_idx >> (8 * i)) & 0xff
        data_frame.append(uint32_package)
    # 6: finish
    data_frame.append(is_finish)
    # 7: b_len
    data_frame.append(end_idx - start_idx)
    # 8-48: data
    for i in range(start_idx, end_idx):
        data_frame.append(buffer[i])
    print("Buffer val is: ", buffer[start_idx: end_idx])

    # send
    # send_data = ''.join(data_frame)
    # return send_data
    return bytes(data_frame)

def send_process(encode_buffer: np.ndarray, outep, inep):
    buffer_len = len(encode_buffer)
    i = 0
    package_idx = 0
    max_data_body = 48 - 7
    while i < buffer_len:
        if buffer_len - i >= max_data_body:
            next_idx = i + max_data_body
            is_finish = 0
        else:
            next_idx = i + (buffer_len - i)
            is_finish = 1

        # make frame
        one_frame = make_one_frame(encode_buffer, package_idx, i, next_idx, is_finish)

        # send try
        status = False
        for i in range(3):
            status = send_and_read(outep, inep, one_frame)
            if not status:
                time.sleep(1)
            else:
                break

        if not status:
            print("Data Send error in {}, data frame is: {}".format(package_idx, one_frame))
            return False

        # increase
        i = next_idx
        package_idx += 1

        print("Debug: package: {}, i_pointer: {}, left: {}".format(package_idx, i, buffer_len - i))

    return True


if __name__ == '__main__':
    # encode image
    image = cv.imread("sample-image.jpg")
    image = cv.resize(image, (20, 20))
    # cv.imwrite("debug.jpg", image)
    # ret, encode_image = cv.imencode('.jpg', image, [cv.IMWRITE_JPEG_QUALITY, 100])
    encode_image = image.flatten()

    # init usb
    outep, inep = init_usb()

    # test send
    start_time = time.perf_counter()
    send_process(encode_image, outep, inep)
    end_time = time.perf_counter()

    cost = end_time - start_time
    data_size = 1 * len(encode_image) / 1e3  # Byte --1e3--> KB --1e3--> MB
    print("time cost in {}, about {}KB/s".format(cost, data_size / cost))
